import cv2
import numpy as np

class H264Player:
    def __init__(self, filename):
        self.cap = cv2.VideoCapture(filename)
        self.frames = []
        self.current_frame_idx = 0
        
        # Read all frames into memory
        while True:
            ret, frame = self.cap.read()
            if not ret:
                break
            self.frames.append(frame)
            
        self.cap.release()
        print(f"Loaded {len(self.frames)} frames")
        
    def show_current_frame(self):
        if 0 <= self.current_frame_idx < len(self.frames):
            cv2.imshow('frame', self.frames[self.current_frame_idx])
            print(f"Showing frame {self.current_frame_idx}/{len(self.frames)-1}")
            
    def next_frame(self):
        if self.current_frame_idx < len(self.frames) - 1:
            self.current_frame_idx += 1
            self.show_current_frame()
            
    def prev_frame(self):
        if self.current_frame_idx > 0:
            self.current_frame_idx -= 1
            self.show_current_frame()

def main():
    import sys
    if len(sys.argv) != 2:
        print("Usage: python readh264.py <h264_file>")
        return
        
    player = H264Player(sys.argv[1])
    
    if len(player.frames) == 0:
        print("No frames found in file")
        return
        
    player.show_current_frame()
    
    while True:
        key = cv2.waitKey(0) & 0xFF
        
        if key == ord('f'):  # Forward
            player.next_frame()
        elif key == ord('b'):  # Backward
            player.prev_frame()
        elif key == ord('q'):  # Quit
            break
            
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
